package com.guardchina.framework.mybatis.datasource;

import com.guardchina.framework.mybatis.api.RoutingDataSourceStrategy;
import com.guardchina.framework.mybatis.enu.DataSourceType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 描述：动态数据源切换
 *
 * @author lwb
 * @date 2019/3/22 9:30
 */
@Slf4j
public class RoutingDataSource extends AbstractRoutingDataSource {

    private String[] slaveDataSourceLookUpKeys;

    private com.guardchina.framework.mybatis.datasource.RoutingDataSourceProvider routingDataSourceProvider;
    private RoutingDataSourceStrategy routingDataSourceStrategy;

    public RoutingDataSource(RoutingDataSourceProvider routingDataSourceProvider, RoutingDataSourceStrategy routingDataSourceStrategy){
        this.routingDataSourceProvider = routingDataSourceProvider;
        this.routingDataSourceStrategy = routingDataSourceStrategy;
    }

    @Override
    protected Object determineCurrentLookupKey() {

        DataSourceType dataSourceType = RoutingDataSourceContextHolder.getDataSourceType();

        //从库数据源
        if(DataSourceType.slave == dataSourceType){
            String slaveDataSourceKey = routingDataSourceStrategy.deterMineSlaveDataSource(slaveDataSourceLookUpKeys);
            log.debug("Switch to slave data source of name {}", slaveDataSourceLookUpKeys);
            return slaveDataSourceKey;
        }

        //主库数据源
        log.debug("use master data source of name {}", dataSourceType.name());
        return dataSourceType.name();
    }

    @Override
    public void afterPropertiesSet() {
        DataSource masterDataSource = routingDataSourceProvider.loadMasterDataSource();
        Map<String, DataSource> slaveDataSources = routingDataSourceProvider.loadSlaveDataSources();

        Set<String> slaveDataSourceKeys = slaveDataSources.keySet();

        this.slaveDataSourceLookUpKeys = slaveDataSourceKeys.toArray(new String[0]);

        Map<Object, Object> targetDataSource = new HashMap<>(slaveDataSources.size() + 1);

        targetDataSource.put(DataSourceType.master.name(), masterDataSource);
        log.info("设置主库{}", masterDataSource);
        targetDataSource.putAll(slaveDataSources);
        log.info("设置从库{}", slaveDataSourceKeys);

        super.setDefaultTargetDataSource(masterDataSource);
        super.setTargetDataSources(targetDataSource);
        super.afterPropertiesSet();
    }
}
